#include "RgCSDAScorer.hh"

#include "G4EmCalculator.hh"
#include "G4Material.hh"
#include "G4ParticleDefinition.hh"
#include "G4Step.hh"
#include "G4Track.hh"
#include "G4UnitsTable.hh"

RgCSDAScorer::RgCSDAScorer(const G4String& name, G4int depth) : G4VPrimitiveScorer(name, depth) {}

//-------------- 每次 Run 开始时清零 --------------
void RgCSDAScorer::Initialize(G4HCofThisEvent* HCE)
{
    const G4String& detName = GetMultiFunctionalDetector()->GetName();
    fCSDAMap = new G4THitsMap<G4double>(detName, GetName());

    if (fHCID < 0) fHCID = GetCollectionID(0);
    HCE->AddHitsCollection(fHCID, fCSDAMap);
    fTrkMap.clear();
}

//-------------- 事件结束（可留空） --------------
void RgCSDAScorer::EndOfEvent(G4HCofThisEvent*) {}

//-------------- 真正干活的地方 --------------
G4bool RgCSDAScorer::ProcessHits(G4Step* step, G4TouchableHistory*)
{
    const G4Track* track = step->GetTrack();
    G4int id = track->GetTrackID();

    // 1. 第一次碰见这条 track：记初始动能
    if (fTrkMap.find(id) == fTrkMap.end())
    {
        fTrkMap[id].ekin0 = track->GetKineticEnergy() / CLHEP::MeV;
    }

    // 2. 粒子停了吗？（动能 = 0 或即将被 Kill）
    G4bool stopped = (track->GetKineticEnergy() == 0.0) || (track->GetTrackStatus() == fStopAndKill);

    if (!stopped) return true;

    // 3. 查 CSDA 表
    G4EmCalculator emCalc;
    // 初始动能; ,track->GetMaterialCutsCouple()->GetIndex()
    G4double csdaRange = emCalc.GetCSDARange(fTrkMap[id].ekin0 * CLHEP::MeV, track->GetParticleDefinition(), track->GetMaterial());

    // 4. 写进 hits map（index=0 表示整个 mesh 唯一 bin）
    G4int idx = 0;
    fCSDAMap->add(idx, csdaRange);

    // 5. 清理，防止重复
    fTrkMap.erase(id);
    return true;
}

//-------------- 打印调试 --------------
void RgCSDAScorer::PrintAll()
{
    G4cout << "-------- CSDA Range (from NIST table) --------" << G4endl;
    auto* m = fCSDAMap->GetMap();
    for (const auto& [idx, val] : *m) G4cout << " idx=" << idx << "  CSDA=" << G4BestUnit(*val, "Length") << G4endl;
}